#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
数据集浏览器组件 - 用于显示和管理数据集中的图像
"""

import os
from PyQt5.QtWidgets import (
    QWidget, QVBoxLayout, QHBoxLayout, QListWidget, QListWidgetItem,
    QPushButton, QLabel, QSplitter, QFileDialog, QMenu, QAction,
    QToolBar, QComboBox, QLineEdit, QMessageBox
)
from PyQt5.QtCore import Qt, QSize, pyqtSignal, QRect
from PyQt5.QtGui import QPixmap, QIcon, QImage

import cv2
import numpy as np

class ImageListWidget(QListWidget):
    """自定义列表组件，用于显示图像缩略图"""
    
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setIconSize(QSize(120, 120))
        self.setViewMode(QListWidget.IconMode)
        self.setResizeMode(QListWidget.Adjust)
        self.setMovement(QListWidget.Static)
        self.setSpacing(5)
        
        # 右键菜单
        self.setContextMenuPolicy(Qt.CustomContextMenu)
        self.customContextMenuRequested.connect(self.show_context_menu)
    
    def show_context_menu(self, pos):
        """显示右键菜单"""
        # 获取选中项
        item = self.itemAt(pos)
        if not item:
            return
        
        # 创建菜单
        menu = QMenu(self)
        
        # 添加菜单项
        view_action = QAction("查看图像", self)
        annotate_action = QAction("标注图像", self)
        delete_action = QAction("删除图像", self)
        
        menu.addAction(view_action)
        menu.addAction(annotate_action)
        menu.addSeparator()
        menu.addAction(delete_action)
        
        # 连接信号
        view_action.triggered.connect(lambda: self.parent().view_image(item))
        annotate_action.triggered.connect(lambda: self.parent().annotate_image(item))
        delete_action.triggered.connect(lambda: self.parent().delete_image(item))
        
        # 显示菜单
        menu.exec_(self.mapToGlobal(pos))


class DatasetBrowser(QWidget):
    """数据集浏览器组件"""
    
    # 自定义信号
    image_selected = pyqtSignal(str)  # 图像被选中，传递图像路径
    image_double_clicked = pyqtSignal(str)  # 图像被双击，传递图像路径
    image_deleted = pyqtSignal(str)  # 图像被删除，传递图像路径
    
    def __init__(self, parent=None):
        """初始化"""
        super().__init__(parent)
        
        # 设置内部变量
        self.current_dataset_path = None
        self.image_paths = []
        self.search_text = ""
        self.current_label_filter = "全部"
        
        # 初始化界面
        self._setup_ui()
    
    def _setup_ui(self):
        """设置界面"""
        # 主布局
        main_layout = QVBoxLayout(self)
        main_layout.setContentsMargins(0, 0, 0, 0)
        
        # 工具栏
        toolbar = QToolBar()
        
        # 添加图像按钮
        self.add_image_btn = QPushButton("添加图像")
        self.add_image_btn.clicked.connect(self.on_add_image)
        toolbar.addWidget(self.add_image_btn)
        
        # 标签过滤器
        toolbar.addSeparator()
        toolbar.addWidget(QLabel("标签过滤:"))
        self.label_filter = QComboBox()
        self.label_filter.addItem("全部")
        self.label_filter.currentTextChanged.connect(self.on_label_filter_changed)
        toolbar.addWidget(self.label_filter)
        
        # 搜索框
        toolbar.addSeparator()
        toolbar.addWidget(QLabel("搜索:"))
        self.search_box = QLineEdit()
        self.search_box.setPlaceholderText("输入文件名搜索")
        self.search_box.textChanged.connect(self.on_search_text_changed)
        toolbar.addWidget(self.search_box)
        
        main_layout.addWidget(toolbar)
        
        # 图像列表
        self.image_list = ImageListWidget()
        self.image_list.itemSelectionChanged.connect(self.on_selection_changed)
        self.image_list.itemDoubleClicked.connect(self.on_item_double_clicked)
        
        main_layout.addWidget(self.image_list)
        
        # 状态标签
        self.status_label = QLabel("就绪")
        main_layout.addWidget(self.status_label)
    
    def set_dataset(self, dataset_path, image_paths=None):
        """
        设置当前数据集
        
        Args:
            dataset_path (str): 数据集路径
            image_paths (list, optional): 图像路径列表
        """
        self.current_dataset_path = dataset_path
        
        # 清除当前列表
        self.image_list.clear()
        self.image_paths = []
        
        if not dataset_path:
            self.status_label.setText("没有打开的数据集")
            return
        
        # 如果没有提供图像路径，尝试从数据集路径中查找
        if not image_paths:
            images_dir = os.path.join(dataset_path, "images")
            if os.path.exists(images_dir):
                # 获取所有图像文件
                image_paths = []
                for filename in os.listdir(images_dir):
                    if filename.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff')):
                        image_paths.append(os.path.join(images_dir, filename))
        
        # 保存图像路径
        self.image_paths = image_paths or []
        
        # 添加图像到列表
        self._populate_image_list()
        
        # 更新状态
        self.status_label.setText(f"共 {len(self.image_paths)} 张图像")
    
    def _populate_image_list(self):
        """填充图像列表"""
        self.image_list.clear()
        
        # 过滤图像
        filtered_paths = self._filter_images()
        
        for path in filtered_paths:
            # 创建缩略图
            thumb = self._create_thumbnail(path, (120, 120))
            
            # 创建列表项
            item = QListWidgetItem()
            item.setIcon(QIcon(thumb))
            item.setText(os.path.basename(path))
            item.setData(Qt.UserRole, path)  # 存储完整路径
            
            # 添加到列表
            self.image_list.addItem(item)
        
        # 更新状态
        self.status_label.setText(f"显示 {self.image_list.count()} / {len(self.image_paths)} 张图像")
    
    def _filter_images(self):
        """
        根据搜索文本和标签过滤图像
        
        Returns:
            list: 过滤后的图像路径列表
        """
        # 如果没有过滤条件，返回所有图像
        if not self.search_text and self.current_label_filter == "全部":
            return self.image_paths
        
        # 根据文件名过滤
        filtered = []
        for path in self.image_paths:
            filename = os.path.basename(path).lower()
            
            # 文件名过滤
            if self.search_text and self.search_text.lower() not in filename:
                continue
            
            # 标签过滤 (需要检查图像的标注)
            if self.current_label_filter != "全部":
                # TODO: 实现标签过滤
                # 这需要访问图像的标注数据，这里暂时不实现
                pass
            
            filtered.append(path)
        
        return filtered
    
    def _create_thumbnail(self, image_path, size):
        """
        创建图像缩略图
        
        Args:
            image_path (str): 图像路径
            size (tuple): 缩略图大小 (width, height)
            
        Returns:
            QPixmap: 缩略图
        """
        try:
            # 读取图像
            img = cv2.imread(image_path)
            if img is None:
                # 创建空白图像
                img = np.zeros((size[1], size[0], 3), dtype=np.uint8)
                cv2.putText(img, "加载失败", (10, size[1]//2), 
                           cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
            
            # 调整大小，保持宽高比
            h, w = img.shape[:2]
            if w > h:
                new_w = size[0]
                new_h = int(h * new_w / w)
            else:
                new_h = size[1]
                new_w = int(w * new_h / h)
            
            # 确保尺寸不为零
            new_w = max(1, new_w)
            new_h = max(1, new_h)
            
            img = cv2.resize(img, (new_w, new_h))
            
            # 创建空白图像
            thumb = np.zeros((size[1], size[0], 3), dtype=np.uint8)
            
            # 将调整后的图像放在中心
            y_offset = (size[1] - new_h) // 2
            x_offset = (size[0] - new_w) // 2
            
            thumb[y_offset:y_offset+new_h, x_offset:x_offset+new_w] = img
            
            # 转换为QPixmap
            rgb_image = cv2.cvtColor(thumb, cv2.COLOR_BGR2RGB)
            h, w, c = rgb_image.shape
            qimg = QImage(rgb_image.data, w, h, w*c, QImage.Format_RGB888)
            return QPixmap(qimg)
            
        except Exception as e:
            print(f"创建缩略图失败: {e}")
            # 返回空白图像
            pixmap = QPixmap(size[0], size[1])
            pixmap.fill(Qt.gray)
            return pixmap
    
    def update_label_filter(self, labels):
        """
        更新标签过滤器选项
        
        Args:
            labels (list): 标签列表
        """
        # 保存当前选中的标签
        current_text = self.label_filter.currentText()
        
        # 清除当前选项
        self.label_filter.clear()
        
        # 添加"全部"选项
        self.label_filter.addItem("全部")
        
        # 添加标签
        for label in labels:
            if isinstance(label, dict) and "name" in label:
                self.label_filter.addItem(label["name"])
            elif isinstance(label, str):
                self.label_filter.addItem(label)
        
        # 恢复选择
        index = self.label_filter.findText(current_text)
        if index >= 0:
            self.label_filter.setCurrentIndex(index)
        else:
            self.label_filter.setCurrentIndex(0)  # 默认选择"全部"
    
    def on_add_image(self):
        """添加图像按钮点击事件"""
        if not self.current_dataset_path:
            QMessageBox.warning(self, "错误", "请先打开数据集")
            return
        
        # 打开文件对话框
        file_dialog = QFileDialog(self)
        file_dialog.setFileMode(QFileDialog.ExistingFiles)
        file_dialog.setNameFilter("图像文件 (*.jpg *.jpeg *.png *.bmp *.tif *.tiff)")
        
        if file_dialog.exec_():
            file_paths = file_dialog.selectedFiles()
            if file_paths:
                # 这里需要调用数据集管理器的添加图像方法
                # 由于这个方法是在父窗口中实现的，我们发射一个自定义信号
                # 在MainWindow中处理这个信号
                # 目前先简单地更新UI
                self.image_paths.extend(file_paths)
                self._populate_image_list()
    
    def on_selection_changed(self):
        """选择改变事件"""
        items = self.image_list.selectedItems()
        if items:
            # 发射信号
            image_path = items[0].data(Qt.UserRole)
            self.image_selected.emit(image_path)
    
    def on_item_double_clicked(self, item):
        """项目双击事件"""
        # 发射信号
        image_path = item.data(Qt.UserRole)
        self.image_double_clicked.emit(image_path)
    
    def on_search_text_changed(self, text):
        """搜索文本改变事件"""
        self.search_text = text
        self._populate_image_list()
    
    def on_label_filter_changed(self, text):
        """标签过滤器改变事件"""
        self.current_label_filter = text
        self._populate_image_list()
    
    def view_image(self, item):
        """查看图像"""
        if item:
            image_path = item.data(Qt.UserRole)
            self.image_double_clicked.emit(image_path)
    
    def annotate_image(self, item):
        """标注图像"""
        if item:
            image_path = item.data(Qt.UserRole)
            self.image_selected.emit(image_path)
            # TODO: 切换到标注模式，这个需要在主窗口中实现
    
    def delete_image(self, item):
        """删除图像"""
        if item:
            image_path = item.data(Qt.UserRole)
            
            # 确认对话框
            reply = QMessageBox.question(
                self, 
                "确认删除", 
                f"确定要删除图像 {os.path.basename(image_path)}？",
                QMessageBox.Yes | QMessageBox.No,
                QMessageBox.No
            )
            
            if reply == QMessageBox.Yes:
                # 从列表中移除
                row = self.image_list.row(item)
                self.image_list.takeItem(row)
                
                # 从路径列表中移除
                if image_path in self.image_paths:
                    self.image_paths.remove(image_path)
                
                # 发射信号
                self.image_deleted.emit(image_path)
                
                # 更新状态
                self.status_label.setText(f"显示 {self.image_list.count()} / {len(self.image_paths)} 张图像")
    
    def refresh(self):
        """刷新列表"""
        self._populate_image_list() 